{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/tools/code-snippet.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jypUkLWc48R_"
      },
      "source": [
        "## Commit trained model to Huggingface"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "182Law9oUiYN",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Clone Model or Datasets\n",
        "\n",
        "#@markdown Type of item to clone (model or dataset)\n",
        "type_of_item = \"dataset\" #@param [\"model\", \"dataset\"]\n",
        "\n",
        "#@markdown Install or uninstall git lfs\n",
        "install_git_lfs = False #@param {'type':'boolean'}\n",
        "\n",
        "%cd /content\n",
        "username = \"Linaqruf\" #@param {'type': 'string'}\n",
        "model_repo = \"your-model-repo\" #@param {'type': 'string'}\n",
        "datasets_repo = \"your-dataset-repo\" #@param {'type': 'string'}\n",
        "\n",
        "if type_of_item == \"model\":\n",
        "  Repository_url = f\"https://huggingface.co/{username}/{model_repo}\"\n",
        "elif type_of_item == \"dataset\":\n",
        "  Repository_url = f\"https://huggingface.co/datasets/{username}/{datasets_repo}\"\n",
        "\n",
        "if install_git_lfs:\n",
        "  !git lfs install\n",
        "else:\n",
        "  !git lfs uninstall\n",
        "!git clone {Repository_url}\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "87wG7QIZbtZE",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title Commit Model or Datasets to Huggingface\n",
        "\n",
        "#@markdown Type of item to commit (model or dataset)\n",
        "type_of_item = \"model\" #@param [\"model\", \"dataset\"]\n",
        "\n",
        "%cd /content\n",
        "#@markdown Go to your model or dataset path\n",
        "item_path = \"your-cloned-repo-name-or-path\" #@param {'type': 'string'}\n",
        "\n",
        "#@markdown #Git Commit\n",
        "\n",
        "#@markdown Set **git commit identity**\n",
        "email = \"your-email\" #@param {'type': 'string'}\n",
        "name = \"your-huggingface-username\" #@param {'type': 'string'}\n",
        "#@markdown Set **commit message**\n",
        "commit_m = \"feat: upload prototype model\" #@param {'type': 'string'}\n",
        "\n",
        "%cd {item_path}\n",
        "!git lfs install\n",
        "!huggingface-cli lfs-enable-largefiles .\n",
        "!git add .\n",
        "!git lfs help smudge\n",
        "!git config --global user.email \"{email}\"\n",
        "!git config --global user.name \"{name}\"\n",
        "!git commit -m \"{commit_m}\"\n",
        "!git push\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training"
      ],
      "metadata": {
        "id": "D2Nr6bw0Ir9s"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Using epochs instead of max training step\n",
        "#@markdown ### Define Parameters\n",
        "\n",
        "import glob\n",
        "import math\n",
        "\n",
        "V2 = \"none\" #@param [\"none\", \"V2_base\", \"V2_768_v\"] {allow-input: false}\n",
        "num_cpu_threads_per_process = 8 #@param {'type':'integer'}\n",
        "save_state = True #@param {'type':'boolean'}\n",
        "train_batch_size = 4  #@param {type: \"slider\", min: 1, max: 10}\n",
        "learning_rate =\"2-e6\" #@param {'type':'string'}\n",
        "num_epoch = 2 #@param {'type':'integer'}\n",
        "dataset_repeats = 1 #@param {'type':'integer'}\n",
        "train_text_encoder = False #@param {'type':'boolean'}\n",
        "lr_scheduler = \"constant\" #@param  [\"linear\", \"cosine\", \"cosine_with_restarts\", \"polynomial\", \"constant\", \"constant_with_warmup\"] {allow-input: false}\n",
        "max_token_length = \"225\" #@param  [\"150\", \"225\"] {allow-input: false}\n",
        "clip_skip = 2 #@param {type: \"slider\", min: 1, max: 10}\n",
        "mixed_precision = \"fp16\" #@param [\"no\",\"fp16\",\"bf16\"] {allow-input: false}\n",
        "save_model_as = \"ckpt\" #@param [\"default\", \"ckpt\", \"safetensors\", \"diffusers\", \"diffusers_safetensors\"] {allow-input: false}\n",
        "save_precision = \"None\" #@param [\"None\",\"float\", \"fp16\", \"bf16\"] {allow-input: false}\n",
        "save_every_n_epochs = 50 #@param {'type':'integer'}\n",
        "gradient_accumulation_steps = 1 #@param {type: \"slider\", min: 1, max: 10}\n",
        "#@markdown ### Log And Debug\n",
        "log_prefix = \"fine-tune-style1\" #@param {'type':'string'}\n",
        "logs_dst = \"/content/fine_tune/training_logs\" #@param {'type':'string'}\n",
        "debug_mode = False #@param {'type':'boolean'}\n",
        "\n",
        "#V2 Inference\n",
        "\n",
        "# Check if directory exists\n",
        "if not os.path.exists(output_dir):\n",
        "  # Create directory if it doesn't exist\n",
        "  os.makedirs(output_dir)\n",
        "\n",
        "inference_url = \"https://raw.githubusercontent.com/Stability-AI/stablediffusion/main/configs/stable-diffusion/\"\n",
        "\n",
        "if V2 == \"V2_base\":\n",
        "  v2_model = \"--v2\"\n",
        "  v2_768v_model= \"\"\n",
        "  inference_url += \"v2-inference.yaml\"\n",
        "elif V2 == \"V2_768_v\":\n",
        "  v2_model = \"--v2\"\n",
        "  v2_768v_model = \"--v2_parameterization\"\n",
        "  inference_url += \"v2-inference-v.yaml\"\n",
        "else:\n",
        "  v2_model = \"\"\n",
        "  v2_768v_model = \"\"\n",
        "\n",
        "try:\n",
        "  if V2 != \"none\":\n",
        "    !wget {inference_url} -O {output_dir}/last.yaml\n",
        "    print(\"File successfully downloaded\")\n",
        "except:\n",
        "  print(\"There was an error downloading the file. Please check the URL and try again.\")\n",
        "\n",
        "if V2 == \"none\":\n",
        "  penultimate_layer = \"--clip_skip\" + \"=\" + \"{}\".format(clip_skip)\n",
        "else:\n",
        "  penultimate_layer = \"\"\n",
        "\n",
        "if save_model_as == \"default\":\n",
        "  sv_model = \"\"\n",
        "else: \n",
        "  sv_model = \"--save_model_as \" + str(save_model_as)\n",
        "\n",
        "if save_state == True:\n",
        "  sv_state = \"--save_state\"\n",
        "else:\n",
        "  sv_state = \"\"\n",
        "\n",
        "if resume_path == \"\":\n",
        "  rs_state = \"\"\n",
        "else:\n",
        "  rs_state = \"--resume \" + str(resume_path)\n",
        "\n",
        "if save_every_n_epochs == 0 :\n",
        "  save_epoch = \"\"\n",
        "else:\n",
        "  save_epoch = \"--save_every_n_epochs\" + \"=\" + \"{}\".format(save_every_n_epochs)\n",
        "\n",
        "if save_precision == \"None\":\n",
        "  sv_precision = \"\"\n",
        "else :\n",
        "  sv_precision = \"--save_precision=\" + str(save_precision)\n",
        "\n",
        "if debug_mode == True:\n",
        "  debug_dataset = \"--debug_dataset\"\n",
        "else:\n",
        "  debug_dataset = \"\"\n",
        "\n",
        "if train_text_encoder == True:\n",
        "  text_encoder = \"--train_text_encoder\"\n",
        "else:\n",
        "  text_encoder = \"\"\n",
        "\n",
        "# Get number of valid images\n",
        "image_num = len(glob.glob(train_data_dir + \"/*.npz\"))\n",
        "\n",
        "print(\"Total Train Data =\", image_num)\n",
        "print(\"Total Epoch=\", num_epoch)\n",
        "print(\"Dataset repeats =\", dataset_repeats, \"x\")\n",
        "repeats = image_num * dataset_repeats\n",
        "print(\"Total Repeats =\", image_num, \"*\", dataset_repeats, \"=\", repeats)\n",
        "\n",
        "# calculate max_train_steps\n",
        "max_train_steps = math.ceil(repeats / train_batch_size * num_epoch)\n",
        "print(\"max_train_steps =\", repeats, \"/\", train_batch_size, \"*\", num_epoch ,\"=\", max_train_steps, \"\\n\")\n",
        "\n",
        "%cd /content/kohya-trainer\n",
        "\n",
        "!accelerate launch \\\n",
        "  --config_file {accelerate_config} \\\n",
        "  --num_cpu_threads_per_process {num_cpu_threads_per_process} \\\n",
        "  fine_tune.py \\\n",
        "  {v2_model} \\\n",
        "  {v2_768v_model} \\\n",
        "  --pretrained_model_name_or_path={pre_trained_model_path} \\\n",
        "  --in_json {meta_lat_json_dir} \\\n",
        "  --train_data_dir={train_data_dir} \\\n",
        "  --output_dir={output_dir} \\\n",
        "  --shuffle_caption \\\n",
        "  --train_batch_size={train_batch_size} \\\n",
        "  --learning_rate={learning_rate} \\\n",
        "  --lr_scheduler={lr_scheduler} \\\n",
        "  --max_token_length={max_token_length} \\\n",
        "  {penultimate_layer} \\\n",
        "  --mixed_precision={mixed_precision} \\\n",
        "  --max_train_steps={max_train_steps} \\\n",
        "  --use_8bit_adam \\\n",
        "  --xformers \\\n",
        "  --gradient_checkpointing \\\n",
        "  --gradient_accumulation_steps {gradient_accumulation_steps} \\\n",
        "  {sv_model} \\\n",
        "  {text_encoder} \\\n",
        "  {sv_state} \\\n",
        "  {rs_state} \\\n",
        "  {save_epoch} \\\n",
        "  {sv_precision} \\\n",
        "  {debug_dataset} \\\n",
        "  --logging_dir={logs_dst} \\\n",
        "  --log_prefix {log_prefix}\n",
        "\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "92jk9Cp3eMHv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "aUiK2tqPBmsp"
      },
      "outputs": [],
      "source": [
        "#@title Auto-move file \n",
        "#@markdown This code automatically moves the `last.ckpt` file, the `last-state` folder, the `train_data`\n",
        "#@markdown directory, and the `meta_lat.json` directory from the output directory to the cloned model\n",
        "#@markdown and datasets repositories. Before running this code, you need to clone the datasets and\n",
        "#@markdown model repositories from huggingface. The code checks for the existence of these files and\n",
        "#@markdown folders in the source and destination directories, and prints messages if they already\n",
        "#@markdown exist or do not exist. It uses the os and shutil libraries to check for the existence of\n",
        "#@markdown files and folders and to move them.\n",
        "import shutil\n",
        "\n",
        "# The path of the output directory\n",
        "output_dir = '/content/kohya-trainer/fine_tuned/' #@param {'type':'string'}\n",
        "\n",
        "# The name of the model\n",
        "model_name = 'momoko30k' #@param {'type':'string'}\n",
        "\n",
        "# The path of the cloned model repository\n",
        "cloned_model_repo = '/content/momoko' #@param {'type':'string'}\n",
        "\n",
        "# The name of the save state\n",
        "save_state_name = 'momoko30k-state' #@param {'type':'string'}\n",
        "\n",
        "# The path of the cloned datasets repository\n",
        "cloned_datasets_repo = '/content/momoko-tag' #@param {'type':'string'}\n",
        "\n",
        "# The path of the meta lat json directory\n",
        "meta_lat_json_dir = \"/content/kohya-trainer/meta_lat.json\" #@param {'type':'string'}\n",
        "\n",
        "# The path of the train data directory\n",
        "train_data_dir = \"/content/kohya-trainer/train_data\" #@param {'type':'string'}\n",
        "\n",
        "if opt_out == True :\n",
        "    # Move file\n",
        "  src_file = f'{output_dir}/last.ckpt'\n",
        "  dst_file = f'{cloned_model_repo}/{model_name}.ckpt'\n",
        "  if os.path.exists(src_file):\n",
        "      if not os.path.exists(dst_file):\n",
        "          shutil.move(src_file, dst_file)\n",
        "          print(f'Moved {src_file} to {dst_file}\\n', flush=True)\n",
        "      else:\n",
        "          print(f'{dst_file} already exists\\n', flush=True)\n",
        "  else:\n",
        "      print(f'There is no {src_file} like that\\n', flush=True)\n",
        "\n",
        "  # Move folder\n",
        "  src_folder = f'{output_dir}/last-state'\n",
        "  dst_folder = f'{cloned_datasets_repo}/{save_state_name}'\n",
        "  if os.path.exists(src_folder):\n",
        "      if not os.path.exists(dst_folder):\n",
        "          shutil.move(src_folder, dst_folder)\n",
        "          print(f'Moved {src_folder} to {dst_folder}\\n', flush=True)\n",
        "      else:\n",
        "          print(f'{dst_folder} already exists\\n', flush=True)\n",
        "  else:\n",
        "      print(f'There is no {src_folder} like that\\n', flush=True)\n",
        "\n",
        "  # Define train data directory\n",
        "  dst_train_data_dir = f'{cloned_datasets_repo}/train_data'\n",
        "\n",
        "  # Check if train data directory already exists\n",
        "  if not os.path.exists(dst_train_data_dir):\n",
        "    # Move train data directory\n",
        "    src_train_data_dir = f'{train_data_dir}'\n",
        "    if os.path.exists(src_train_data_dir):\n",
        "        shutil.move(src_train_data_dir, dst_train_data_dir)\n",
        "        print(f'Moved {src_train_data_dir} to {dst_train_data_dir}\\n', flush=True)\n",
        "    else:\n",
        "        print(f'There is no {src_train_data_dir} like that\\n', flush=True)\n",
        "\n",
        "  # Define meta lat json directory\n",
        "  dst_meta_lat_json_dir = f'{cloned_datasets_repo}/meta_lat.json'\n",
        "\n",
        "  # Check if meta lat json directory already exists\n",
        "  if not os.path.exists(dst_meta_lat_json_dir):\n",
        "    # Move meta lat json directory\n",
        "    src_meta_lat_json_dir = f'{meta_lat_json_dir}'\n",
        "    if os.path.exists(src_meta_lat_json_dir):\n",
        "        shutil.move(src_meta_lat_json_dir, dst_meta_lat_json_dir)\n",
        "        print(f'Moved {src_meta_lat_json_dir} to {dst_meta_lat_json_dir}\\n', flush=True)\n",
        "    else:\n",
        "        print(f'There is no {src_meta_lat_json_dir} like that\\n', flush=True)\n",
        "\n",
        "  # Iterate over all files and folders in the cloned_datasets_repo directory\n",
        "  for filename in os.listdir(cloned_datasets_repo):\n",
        "    # Check if the file or folder is not the save_state_name folder, the train_data folder, or the meta_lat.json file\n",
        "    if filename != save_state_name and filename != os.path.basename(dst_train_data_dir) and filename != os.path.basename(dst_meta_lat_json_dir):\n",
        "      # Get the path of the file or folder\n",
        "      file_path = os.path.join(cloned_datasets_repo, filename)\n",
        "\n",
        "      # Check if the file or folder is a directory (i.e., a folder)\n",
        "      if os.path.isdir(file_path):\n",
        "        # Delete the folder\n",
        "        shutil.rmtree(file_path)\n",
        "        print(f'Deleted folder: {filename}')\n",
        "      else:\n",
        "        # Delete the file\n",
        "        os.remove(file_path)\n",
        "        print(f'Deleted file: {filename}')\n",
        "\n"
      ]
    }
  ]
}